# coding: utf-8
from .op_basic import as_op
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np


def _use_x_y_data(x, y, data):
    if isinstance(x, str):
        assert isinstance(y, str)
        return data[x], data[y]
    else:
        assert data is None
        return x, y


@as_op('sort_scatter')
def sort_scatter(series, ascending=True, size=None):
    if size is not None:
        plt.figure(figsize=size)
    plt.scatter(range(len(series)), series.sort_values(ascending=ascending))

@as_op('conc_vs_conc_barplot')
def conc_vs_conc_barplot(x, y, data, use_rate=False, size=None):
    '''
    feature离散、目标离散。
    feature在x轴上，相当于不同的feature，看每个target类是怎么分布的。
    use_rate=False，直接数数目
    use_rate=True，在每一个feature内，只考虑每个target类的占比
    size : (w, h)
    '''
    if size is not None:
        plt.figure(figsize=size)
    if use_rate:
        count = data.groupby([x, y]).apply(lambda x: len(x))
        group_sum = count.groupby(level=0).sum()
        count = count / group_sum
        count = pd.DataFrame(count, columns=['rate']).reset_index()
        sns.barplot(x=x, y='rate', hue=y, data=count)
    else:
        sns.countplot(x=x, hue=y, data=data)


@as_op('conc_value_counts_plot')
def conc_value_counts_plot(col, data, size=None):
    '''
    对于一个离散型特征或目标。画出它的value_counts()的图。这个图有两个，一个是原始count，一个是取了log的形式。
    可以看这个离散特征内部是不是平衡的。由此看看可不可以构建特征。
    '''
    value_count = data[col].value_counts()
    if size is not None:
        plt.figure(figsize=size)
    else:
        plt.figure()
    value_count.plot(title='value count')
    if size is not None:
        plt.figure(figsize=size)
    else:
        plt.figure()
    np.log(value_count).plot(title='log value count')


@as_op('cont_vs_conc_distplot')
def cont_vs_conc_distplot(x, y, data, hist=False, rug=False, size=(10, 6)):
    '''
    画分布图。feature连续、目标离散。
    size: (w, h)
    '''
    height = size[1]
    ratio = float(size[0]) / size[1]
    g = sns.FacetGrid(data, hue=y, size=height, aspect=ratio)
    g.map(sns.distplot, x, hist=hist, rug=rug)
    g.add_legend()
    plt.show()


@as_op('cont_vs_conc_stripplot')
def cont_vs_conc_stripplot(x, y, data, point_size=None, size=None, jitter=0.0):
    '''
    对每一类画stripplot。feature连续、目标离散。
    point_size为每个点的大小
    size: (w, h)
    '''
    if size is not None:
        plt.figure(figsize=size)
    if point_size is not None:
        sns.stripplot(x=y, y=x, data=data, size=point_size, jitter=jitter)
    else:
        sns.stripplot(x=y, y=x, data=data, jitter=jitter)


@as_op('cont_vs_conc_swarmplot')
def cont_vs_conc_swarmplot(x, y, data, point_size=None, size=None):
    '''
    注意：数据量大的时候这个函数非常慢！
    对每一类画swarmplot。feature连续、目标离散。
    point_size为每个点的大小
    size: (w, h)
    '''
    if size is not None:
        plt.figure(figsize=size)
    if point_size is not None:
        sns.swarmplot(x=y, y=x, data=data, size=point_size)
    else:
        sns.swarmplot(x=y, y=x, data=data)


@as_op('cont_vs_conc_violinplot')
def cont_vs_conc_violinplot(x, y, data, size=None):
    '''
    对每一类画violinplot。feature连续、目标离散。
    size: (w, h)
    '''
    if size is not None:
        plt.figure(figsize=size)
    sns.violinplot(x=y, y=x, data=data)


@as_op('cont_vs_conc_sort_scatter')
def cont_vs_conc_sort_scatter(x, y, df, size=(10, 6), ascending=True):
    '''
    将连续的变量df[x]的值排序，在图中依次画出散点。不同的y画成不同的颜色。
    每一类y横坐标是0~1标准化的index，纵坐标是df[x]值。
    '''
    def myscatter(_, x, label, color):
        '''
        辅助函数，第一个无效。直接从画第二个函数的scatter
        '''
        plt.scatter(np.array(range(len(x))) / len(x), x.sort_values(ascending=ascending), label=label, color=color)

    height = size[1]
    ratio = float(size[0]) / size[1]

    assert x != 'sort_index_normalized'
    assert y != 'sort_index_normalized'
    df = df[[x, y]].copy()
    df['sort_index_normalized'] = 0  # 为了显示x坐标的label。没有任何作用
    g = sns.FacetGrid(df, hue=y, size=height, aspect=ratio)
    g.map(myscatter, "sort_index_normalized", x)
    g.add_legend()
    plt.show()


@as_op('conc_vs_cont_distplot')
def conc_vs_cont_distplot(x, y, data, hist=False, rug=False, size=(10, 6)):
    height = size[1]
    ratio = float(size[0]) / size[1]
    g = sns.FacetGrid(data, hue=x, size=height, aspect=ratio)
    g.map(sns.distplot, y, hist=hist, rug=rug)
    g.add_legend()
    plt.show()


@as_op('conc_vs_cont_stripplot')
def conc_vs_cont_stripplot(x, y, data, point_size=None, size=None, jitter=0.0):
    if size is not None:
        plt.figure(figsize=size)
    if point_size is not None:
        sns.stripplot(x=x, y=y, data=data, size=point_size, jitter=jitter)
    else:
        sns.stripplot(x=x, y=y, data=data, jitter=jitter)


@as_op('conc_vs_cont_swarmplot')
def conc_vs_cont_swarmplot(x, y, data, point_size=None, size=None):
    if size is not None:
        plt.figure(figsize=size)
    if point_size is not None:
        sns.swarmplot(x=x, y=y, data=data, size=point_size)
    else:
        sns.swarmplot(x=x, y=y, data=data)


@as_op('conc_vs_cont_violinplot')
def conc_vs_cont_violinplot(x, y, data, size=None):
    if size is not None:
        plt.figure(figsize=size)
    sns.violinplot(x=x, y=y, data=data)


@as_op('cont_vs_cont_regplot')
def cont_vs_cont_regplot(x, y, data, fit_reg=True):
    '''
    feature连续，目标也连续，画散点图看一看有没有线性关系。
    可以使用fit_reg决定显不显示那条线
    '''
    sns.regplot(x=x, y=y, data=data, fit_reg=fit_reg)
